#include "LocalProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"

void LocalProcess::OnInit()
{
	mLocalEntitys = mGameScene->GetEntitysByType<LocalEntity, Model>();

	std::vector<Ubo> ubos;
	ubos.resize(mLocalEntitys.size());
	for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mLocalEntitys[i]->GetTransform().GetWorldMatrix();
	}

	ENQUEUE_RENDER_COMMAND(BuildModel)([this, ubos](VkCommandBuffer commandBuffer)
		{
			std::vector<Ubo> entityUbos = ubos;

			for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
			{
				Model& model = mLocalEntitys[i]->GetComponent<Model>();

				std::vector<LocalVertex> tempVertexs;
				tempVertexs.resize(model.mVertexs.size());
				for (uint32 j = 0; j < model.mVertexs.size(); ++j)
				{
					tempVertexs[j].position = model.mVertexs[j].Position;
					tempVertexs[j].normal = model.mVertexs[j].Normal;
				}

				mLocalEntitys[i]->mVertexBuffer = VertexBuffer::Builder()
					.SetData(tempVertexs.data())
					.SetSize(tempVertexs.size() * sizeof(LocalVertex))
					.Build();

				mLocalEntitys[i]->mIndexBuffer = IndexBuffer::Builder()
					.SetData(model.mIndices.data())
					.SetSize(model.mIndices.size() * sizeof(uint32))
					.SetIndexCount(model.mIndices.size())
					.Build();

				mLocalEntitys[i]->mUboBuffer = CustomBuffer::Builder()
					.SetData(&entityUbos[i])
					.SetSize(sizeof(Ubo))
					.SetMemoryProperty(VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT)
					.Build();
			}

			mDescriptor = RHIDescriptor::Builder()
				.PushShader<LocalVertexShader>()
				.PushShader<LocalFragmentShader>()
				.SetDescriptorCount(0, mLocalEntitys.size())
				.Build();

			for (auto& entity : mLocalEntitys)
			{
				mDescriptor->SetValue("ubo", entity->mUboBuffer);
				entity->mDescriptorMask = mDescriptor->AllocateDescriptor();
			}

			mColorRenderTarget = ColorRenderTarget::Builder()
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetMultiple(true)
				.Build(this, "ColorRenderTarget");

			mDepthRenderTarget = DepthRenderTarget::Builder()
				.SetMultiple(true)
				.Build(this, "DepthRenderTarget");

			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorRenderTarget)
				.SetDepthRenderTarget(mDepthRenderTarget)
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetCullMode(VK_CULL_MODE_NONE)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorRenderTarget)
				.AddRenderTarget(mDepthRenderTarget)
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build(this, "RenderPass");

			mFrameBuffer = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void LocalProcess::OnTick(float deltaTime)
{
	std::vector<Ubo> ubos;
	ubos.resize(mLocalEntitys.size());
	for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
	{
		ubos[i].proj = mGameScene->GetCamera().GetProj();
		ubos[i].view = mGameScene->GetCamera().GetView();
		ubos[i].model = mLocalEntitys[i]->GetTransform().GetLocalMatrix();
	}

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this, ubos](VkCommandBuffer commandBuffer)
		{
			std::vector<Ubo> entityUbos = ubos;

			for (uint32 i = 0; i < mLocalEntitys.size(); ++i)
			{
				mLocalEntitys[i]->mUboBuffer->WriteMemory(&entityUbos[i]);
			}

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffer->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			for (const auto& entity : mLocalEntitys)
			{
				vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &entity->mDescriptorMask.GetDescriptorSet(), 0, nullptr);

				VkBuffer vertexBuffer = entity->mVertexBuffer->GetBuffer();
				VkDeviceSize offsets = 0;
				vkCmdBindVertexBuffers(commandBuffer, 0, 1, &vertexBuffer, &offsets);

				VkBuffer indexBuffer = entity->mIndexBuffer->GetBuffer();
				vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);

				vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(glm::vec4), &entity->mMaterial.albedo);

				vkCmdDrawIndexed(commandBuffer, entity->mIndexBuffer->GetIndexCount(), 1, 0, 0, 0);
			}

			vkCmdEndRenderPass(commandBuffer);
		});
}

void LocalProcess::OnResize()
{
	mColorRenderTarget = ColorRenderTarget::Builder()
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetMultiple(true)
		.Build(this, "ColorRenderTarget");

	mDepthRenderTarget = DepthRenderTarget::Builder()
		.SetMultiple(true)
		.Build(this, "DepthRenderTarget");

	mFrameBuffer = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorRenderTarget)
		.AddRenderTarget(mDepthRenderTarget)
		.SetRenderPass(mRenderPass)
		.Build();
}
